package com.element.security.server.web;

import com.common.core.constant.SecurityHeader;
import com.common.core.results.SingleResult;
import com.common.core.results.VoidResult;
import com.element.security.exception.AuthException;
import com.element.security.model.co.TokenCO;
import com.element.security.model.co.TokenRefreshCO;
import com.element.security.server.global.GlobalOAuth2RequestValidator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.ResponseEntity;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.security.oauth2.provider.ClientRegistrationException;
import org.springframework.security.oauth2.provider.endpoint.TokenEndpoint;
import org.springframework.web.HttpRequestMethodNotSupportedException;
import org.springframework.web.bind.annotation.*;

import javax.annotation.Resource;
import java.security.Principal;
import java.util.Map;

/**
 * @auther WuJun.Zhang
 * @date 2022/2/18 下午10:32
 */

@RestController
@RequestMapping(value = "oauth")
public class AuthTokenController {

    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    @Resource
    private TokenEndpoint tokenEndpoint;

    @Resource
    private GlobalOAuth2RequestValidator globalOAuth2RequestValidator;

    @GetMapping(value = "token")
    public SingleResult<TokenCO> getAccessToken(Principal principal, @RequestParam Map<String, String> parameters) throws HttpRequestMethodNotSupportedException {
        tokenEndpoint.setOAuth2RequestValidator(globalOAuth2RequestValidator);
        ResponseEntity<OAuth2AccessToken> responseEntity = tokenEndpoint.getAccessToken(principal, parameters);
        return SingleResult.of(getToken(responseEntity));
    }

    @PostMapping(value = "token")
    public SingleResult<TokenCO> postAccessToken(Principal principal, @RequestParam Map<String, String> parameters) throws HttpRequestMethodNotSupportedException {
        tokenEndpoint.setOAuth2RequestValidator(globalOAuth2RequestValidator);
        ResponseEntity<OAuth2AccessToken> responseEntity = tokenEndpoint.postAccessToken(principal, parameters);
        return SingleResult.of(getToken(responseEntity));
    }

    @ExceptionHandler({HttpRequestMethodNotSupportedException.class})
    public VoidResult handleHttpRequestMethodNotSupportedException(HttpRequestMethodNotSupportedException e) throws Exception {
        logger.error("HttpRequestMethodNotSupportedException:", e);
        ResponseEntity<OAuth2Exception> responseEntity = tokenEndpoint.handleHttpRequestMethodNotSupportedException(e);
        AuthException authException = (AuthException) responseEntity.getBody();
        assert authException != null;
        return authException.getResult();
    }

    @ExceptionHandler({Exception.class})
    public VoidResult handleException(Exception e) throws Exception {
        logger.error("Exception:", e);
        ResponseEntity<OAuth2Exception> responseEntity = tokenEndpoint.handleException(e);
        AuthException authException = (AuthException) responseEntity.getBody();
        assert authException != null;
        return authException.getResult();
    }

    @ExceptionHandler({ClientRegistrationException.class})
    public VoidResult handleClientRegistrationException(ClientRegistrationException e) throws Exception {
        logger.error("ClientRegistrationException:", e);
        ResponseEntity<OAuth2Exception> responseEntity = tokenEndpoint.handleClientRegistrationException(e);
        AuthException authException = (AuthException) responseEntity.getBody();
        assert authException != null;
        return authException.getResult();
    }

    @ExceptionHandler({OAuth2Exception.class})
    public VoidResult handleException(OAuth2Exception e) throws Exception {
        logger.error("OAuth2Exception:", e);
        ResponseEntity<OAuth2Exception> responseEntity = tokenEndpoint.handleException(e);
        AuthException authException = (AuthException) responseEntity.getBody();
        assert authException != null;
        return authException.getResult();
    }

    private TokenCO getToken(ResponseEntity<OAuth2AccessToken> responseEntity) {
        if (null != responseEntity) {
            OAuth2AccessToken oAuth2AccessToken = responseEntity.getBody();
            if (null != oAuth2AccessToken) {
                OAuth2RefreshToken oAuth2RefreshToken = oAuth2AccessToken.getRefreshToken();
                TokenCO token;
                if (null == oAuth2RefreshToken) {
                    token = new TokenCO();
                } else {
                    token = new TokenRefreshCO();
                    ((TokenRefreshCO) token).setRefreshToken(oAuth2RefreshToken.getValue());
                }
                token.setAccessToken(oAuth2AccessToken.getValue());
                token.setExpiresIn(oAuth2AccessToken.getExpiresIn());
                // TODO 返回用户请求头值和加盐
                token.setTokenHeader(SecurityHeader.AUTH_SECURITY_HEADERS);
                token.setTokenType(SecurityHeader.AUTH_SECURITY_VALUE);
                return token;
            }
        }
        return null;
    }
}
